{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "99a5206b-b112-401b-8815-86b3a217d085",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 导入 AKShare 库，用于获取股票数据\n",
    "import akshare as ak\n",
    "# 导入 Pandas 库，用于数据处理和分析\n",
    "import pandas as pd\n",
    "# 导入 datetime 用于处理日期时间\n",
    "from datetime import datetime\n",
    "# 导入 relativedelta 用于计算相对时间差\n",
    "from dateutil.relativedelta import relativedelta\n",
    "# 导入 TA-Lib 库，用于计算技术指标\n",
    "import talib\n",
    "# 导入 MinMaxScaler 用于数据标准化\n",
    "from sklearn.preprocessing import MinMaxScaler\n",
    "# 导入 Keras 的 Sequential 模型\n",
    "from tensorflow.keras.models import Sequential\n",
    "# 导入 LSTM 和 Dense 层用于构建神经网络\n",
    "from tensorflow.keras.layers import LSTM, Dense\n",
    "# 导入 NumPy 库，用于数组操作\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0e0a6e85-2fdb-4229-bed3-24d883e85d3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_stock():\n",
    "    \"\"\"\n",
    "    \t获取 A 股实时行情数据，仅保留股票代码和名称两列\n",
    "    \"\"\"\n",
    "    return ak.stock_zh_a_spot_em()[['代码', '名称']]\n",
    "    \n",
    "def preprocessing(data):\n",
    "    \"\"\"\n",
    "        数据预处理：清洗数据并计算技术指标\n",
    "        :param data: 包含股票历史数据的 DataFrame\n",
    "        :return: 添加了技术指标的 DataFrame\n",
    "    \"\"\"\n",
    "    # 计算 5 日简单移动平均线 (SMA)\n",
    "    data['MA5'] = talib.SMA(data['收盘'], timeperiod=5)\n",
    "    # 计算 MACD 指标，包括 MACD 线、信号线和柱状图\n",
    "    data['MACD'], data['Signal'], data['Hist'] = talib.MACD(data['收盘'], fastperiod=12, slowperiod=26, signalperiod=9)\n",
    "    # 输出最近 5 行的关键数据，便于检查结果\n",
    "    # print(f'----------------查看结果: ------------\\n{data[[\"日期\", \"收盘\", \"MA5\", \"MACD\", \"Signal\", \"Hist\"]].tail()}')\n",
    "    # 返回处理后的数据\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2d056872-fa72-4300-b470-0ef3362672c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trend_analysis(data, code, months):\n",
    "    \"\"\"\n",
    "        进行股票历史趋势分析，动态默认获取最近三个月的数据\n",
    "        :param data: 包含股票代码的 DataFrame\n",
    "        :param code: 具体的股票代码\n",
    "        :return: 包含历史数据的 DataFrame\n",
    "    \"\"\"\n",
    "    # 获取当前日期作为结束日期\n",
    "    end_date = datetime.today()\n",
    "    # 计算三个月前的开始日期\n",
    "    start_date = end_date - relativedelta(months=months)\n",
    "    # 格式化结束日期为 YYYYMMDD\n",
    "    end_date_str = end_date.strftime('%Y%m%d')\n",
    "    # 格式化开始日期为 YYYYMMDD\n",
    "    start_date_str = start_date.strftime('%Y%m%d')\n",
    "    # 调用 AKShare 获取指定股票的日线历史数据\n",
    "    result = ak.stock_zh_a_hist(symbol=code, period='daily', start_date=start_date_str, end_date=end_date_str)\n",
    "    # 输出获取到的数据行数\n",
    "    print(f\"股票 {code} 的历史数据行数: {len(result)}\")\n",
    "    # 返回历史数据\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f0f95593-1ec1-43a6-b457-6b316d5ea4b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fuzzy_match_df(df, query):\n",
    "    \"\"\"\n",
    "        从 DataFrame 中模糊匹配股票代码或名称\n",
    "        :param df: 输入的 DataFrame，包含股票代码和名称\n",
    "        :param query: 用户输入的查询字符串，用于模糊匹配\n",
    "        :return: 匹配到的股票数据的 DataFrame 子集\n",
    "    \"\"\"\n",
    "    # 使用布尔掩码筛选包含查询字符串的行，na=False 避免 NaN 报错\n",
    "    mask = (df['代码'].str.contains(query, na=False) | \n",
    "            df['名称'].str.contains(query, na=False))\n",
    "    # 输出匹配到的数据条数\n",
    "    print(f'匹配到了：{len(df[mask])}条数据')\n",
    "    # 返回匹配结果\n",
    "    return df[mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4cb3d9a5-a114-4133-9531-7d418bb4fbb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trend_analysis(data, code, months):\n",
    "    \"\"\"\n",
    "        进行股票历史趋势分析，动态默认获取最近三个月的数据\n",
    "        :param data: 包含股票代码的 DataFrame\n",
    "        :param code: 具体的股票代码\n",
    "        :return: 包含历史数据的 DataFrame\n",
    "    \"\"\"\n",
    "    # 获取当前日期作为结束日期\n",
    "    end_date = datetime.today()\n",
    "    # 计算三个月前的开始日期\n",
    "    start_date = end_date - relativedelta(months=months)\n",
    "    # 格式化结束日期为 YYYYMMDD\n",
    "    end_date_str = end_date.strftime('%Y%m%d')\n",
    "    # 格式化开始日期为 YYYYMMDD\n",
    "    start_date_str = start_date.strftime('%Y%m%d')\n",
    "    # 调用 AKShare 获取指定股票的日线历史数据\n",
    "    result = ak.stock_zh_a_hist(symbol=code, period='daily', start_date=start_date_str, end_date=end_date_str)\n",
    "    # 输出获取到的数据行数\n",
    "    print(f\"股票 {code} 的历史数据行数: {len(result)}\")\n",
    "    # 返回历史数据\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5879811f-fb60-49f9-8d67-aa72060dbcc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def trend_forecast(data):\n",
    "    \"\"\"\n",
    "        趋势预测：使用 LSTM 模型预测下一交易日的收盘价\n",
    "        :param data: 包含历史数据和技术指标的 DataFrame\n",
    "        :return: 预测的收盘价，或 None（数据不足时）\n",
    "    \"\"\"\n",
    "    # 初始化 MinMaxScaler 用于将数据缩放到 [0, 1]\n",
    "    scaler = MinMaxScaler()\n",
    "    # 将收盘价转换为二维数组并进行标准化\n",
    "    scaled_data = scaler.fit_transform(data['收盘'].values.reshape(-1, 1))\n",
    "    # 定义回溯窗口大小为 60 天\n",
    "    look_back = 60\n",
    "    \n",
    "    # 检查数据长度是否足够进行预测\n",
    "    if len(scaled_data) <= look_back:\n",
    "        # 数据不足时输出提示并返回 None\n",
    "        print(f\"数据长度 {len(scaled_data)} 不足 {look_back}，无法预测\")\n",
    "        return None\n",
    "    \n",
    "    def create_dataset(dataset, look_back=60):\n",
    "        \"\"\"\n",
    "            创建适用于 LSTM 的数据集\n",
    "            :param dataset: 标准化的数据\n",
    "            :param look_back: 回溯窗口大小\n",
    "            :return: 输入特征 X 和目标值 y\n",
    "        \"\"\"\n",
    "        # 初始化输入和输出列表\n",
    "        X, y = [], []\n",
    "        # 遍历数据，生成特征和目标值\n",
    "        for i in range(len(dataset) - look_back):\n",
    "            # 添加前 look_back 天的数据作为特征\n",
    "            X.append(dataset[i:i + look_back])\n",
    "            # 添加第 look_back + 1 天的值作为目标\n",
    "            y.append(dataset[i + look_back])\n",
    "        # 转换为 NumPy 数组\n",
    "        return np.array(X), np.array(y)\n",
    "    \n",
    "    # 创建训练数据集\n",
    "    X, y = create_dataset(scaled_data)\n",
    "    # 调整 X 的形状为 (样本数, 时间步, 特征数)\n",
    "    X = np.reshape(X, (X.shape[0], X.shape[1], 1))\n",
    "    \n",
    "    # 构建 LSTM 模型\n",
    "    model = Sequential()\n",
    "    # 添加第一层 LSTM，50 个单元，返回序列以连接下一层\n",
    "    model.add(LSTM(50, return_sequences=True, input_shape=(look_back, 1)))\n",
    "    # 添加第二层 LSTM，50 个单元\n",
    "    model.add(LSTM(50))\n",
    "    # 添加全连接层，输出 1 个预测值\n",
    "    model.add(Dense(1))\n",
    "    # 编译模型，使用 Adam 优化器和均方误差损失函数\n",
    "    model.compile(optimizer='adam', loss='mean_squared_error')\n",
    "    \n",
    "    # 训练模型，10 个 epoch，批次大小 32\n",
    "    model.fit(X, y, epochs=10, batch_size=32)\n",
    "    \n",
    "    # 使用最后 60 天数据预测下一交易日\n",
    "    predicted = model.predict(X[-1].reshape(1, look_back, 1))\n",
    "    # 反标准化预测值，恢复到原始价格范围\n",
    "    predicted_price = scaler.inverse_transform(predicted)\n",
    "    # 输出预测结果\n",
    "    print(f\"预测下一交易日收盘价: {predicted_price[0][0]}\")\n",
    "    # 返回预测价格\n",
    "    return predicted_price"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b75177b4-456b-427d-a6de-ddd57032c1b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main_implement(query,months=3):\n",
    "    \"\"\"\n",
    "        主函数：执行模糊匹配、趋势分析和预测\n",
    "        :param query: 用户输入的查询字符串\n",
    "        :param months: 查询最近几月份默认3个月\n",
    "    \"\"\"\n",
    "    # 模糊匹配股票\n",
    "    result = fuzzy_match_df(get_stock(), query)\n",
    "    # 如果匹配到数据\n",
    "    if not result.empty:\n",
    "        # 遍历前 2 个匹配结果\n",
    "        for code, name in zip(result['代码'], result['名称']):\n",
    "            # 输出正在处理的股票信息\n",
    "            print(f\"处理股票: {code} - {name}\")\n",
    "            # 获取历史数据\n",
    "            trend_analysis_result = trend_analysis(result, code,months)\n",
    "            # 如果历史数据不为空\n",
    "            if not trend_analysis_result.empty:\n",
    "                # 预处理数据\n",
    "                preprocessing_result = preprocessing(trend_analysis_result)\n",
    "                # 进行趋势预测\n",
    "                trend_forecast(preprocessing_result)\n",
    "            # 如果没有历史数据\n",
    "            else:\n",
    "                # 输出无历史数据提示\n",
    "                print(f\"股票 {code} 无历史数据\")\n",
    "    # 如果未匹配到数据\n",
    "    else:\n",
    "        # 输出未匹配到数据的提示\n",
    "        print('没有匹配到你想要的数据')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1b390dce-6422-4a90-bfe8-37f038f2244b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "匹配到了：1条数据\n",
      "处理股票: 300085 - 银之杰\n",
      "股票 300085 的历史数据行数: 118\n",
      "Epoch 1/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Software\\Anaconda3\\envs\\py_env_3\\Lib\\site-packages\\keras\\src\\layers\\rnn\\rnn.py:200: UserWarning: Do not pass an `input_shape`/`input_dim` argument to a layer. When using Sequential models, prefer using an `Input(shape)` object as the first layer in the model instead.\n",
      "  super().__init__(**kwargs)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m2s\u001b[0m 40ms/step - loss: 0.3525\n",
      "Epoch 2/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 41ms/step - loss: 0.2260\n",
      "Epoch 3/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 37ms/step - loss: 0.1190\n",
      "Epoch 4/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 35ms/step - loss: 0.0366\n",
      "Epoch 5/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 45ms/step - loss: 0.0147\n",
      "Epoch 6/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 45ms/step - loss: 0.0473\n",
      "Epoch 7/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 45ms/step - loss: 0.0277\n",
      "Epoch 8/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 40ms/step - loss: 0.0091\n",
      "Epoch 9/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 33ms/step - loss: 0.0093\n",
      "Epoch 10/10\n",
      "\u001b[1m2/2\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 40ms/step - loss: 0.0145\n",
      "\u001b[1m1/1\u001b[0m \u001b[32m━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[37m\u001b[0m \u001b[1m0s\u001b[0m 245ms/step\n",
      "预测下一交易日收盘价: 33.510128021240234\n"
     ]
    }
   ],
   "source": [
    "# 执行函数，查询包含“关键字或者代码”的股票\n",
    "main_implement('银之杰',6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36a835ea-a373-436d-9d4d-4c76e9ab5a41",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
